
-- set search_path to sm_sc;
-- drop function if exists sm_sc.fv_d_redistr_zscore(float[], float[]);
create or replace function sm_sc.fv_d_redistr_zscore
(
  i_depdt_var                  float[]                 ,                     -- zscore 输出
  -- i_dloss_ddepdt           float[]                 ,                     -- 此入参传入 dloss/dindepdt, 用于 zscore 直接求取 dloss/ddepdt
  i_indepdt_var                  float[]                                       -- zscore 算子的输入，来自上一层算子的输出
)
returns float[]     -- 输出列序与 i_indepdt_var 枚举值 one_hot 序一致
as
$$
declare 
  v_len               int;
  v_stddev_samp       float;

begin
  -- 审计维度
  if current_setting('pg4ml._v_is_debug_check', true) = '1'
  then
    if array_ndims(i_depdt_var) > 4 or array_ndims(i_indepdt_var) > 4
      or array_dims(i_depdt_var) <> array_dims(i_indepdt_var) 
      -- -- or array_length(i_depdt_var, 1) <> array_length(i_indepdt_var, 1)
      -- -- or array_length(i_depdt_var, 2) <> array_length(i_indepdt_var, 2)
      -- -- or array_length(i_depdt_var, 3) <> array_length(i_indepdt_var, 3)
      -- -- or array_length(i_depdt_var, 4) <> array_length(i_indepdt_var, 4)
    then 
      raise exception 'unmatch ndims or length between i_depdt_var and i_indepdt_var.';
    end if;
  end if;

  if i_depdt_var is null
  then
    i_depdt_var := sm_sc.fv_redistr_zscore(i_indepdt_var);
  end if;
  
  v_stddev_samp := sm_sc.fv_aggr_slice_stddev_samp(i_indepdt_var);
  v_len := array_length(i_depdt_var, 1) * coalesce(array_length(i_depdt_var, 2), 1) * coalesce(array_length(i_depdt_var, 3), 1) * coalesce(array_length(i_depdt_var, 4), 1);

  if abs(v_stddev_samp) < 1e-256
  then 
    return array_fill(1.0 :: float, (select array_agg(array_length(i_indepdt_var, a_dim) order by a_dim) from generate_series(1, array_ndims(i_indepdt_var)) tb_a_dim(a_dim)));
  else
    return 
      -- (
      (
        (
          (- 1.0 :: float/ (v_len * v_stddev_samp)) 
          *` (i_depdt_var *` (|@+| i_depdt_var))
        ) 
        -` (1.0 :: float/ (v_stddev_samp * v_len))
      )
      +` (i_depdt_var +` ((v_len - 1.0 :: float) / v_stddev_samp))
      -- )  
      -- *` i_dloss_ddepdt
    ; -- ~=` 8
  end if;

end
$$
language plpgsql stable
parallel safe
cost 100;
-- -- set search_path to sm_sc;
-- select sm_sc.fv_d_redistr_zscore
--   (
--     sm_sc.fv_redistr_zscore(array[[1, 2, 3, 4, 5]]),
--     array[[1, 2, 3, 4, 5]]
--   );

-- select sm_sc.fv_d_redistr_zscore
--   (
--     sm_sc.fv_d_redistr_zscore(array[array[1], array[2], array[3], array[4], array[5]]),
--     -- -` array[array[0.0 :: float], array[0.0 :: float], array[0.0 :: float], array[0.0 :: float], array[1.0 :: float]] /` sm_sc.fv_standlize_mx_zscore(array[array[1], array[2], array[3], array[4], array[5]]),
--     array[array[1], array[2], array[3], array[4], array[5]]
--   );

-- with 
-- cte_arr as 
-- (
--   select 
--     sm_sc.fv_new_rand(array[2, 3, 4]) as a_arr
-- )
-- select 
--   sm_sc.fv_d_redistr_zscore(sm_sc.fv_redistr_zscore(a_arr), a_arr) 
-- from cte_arr

-- with 
-- cte_arr as 
-- (
--   select 
--     sm_sc.fv_new_rand(array[2, 3, 4, 3]) as a_arr
-- )
-- select 
--   sm_sc.fv_d_redistr_zscore(sm_sc.fv_redistr_zscore(a_arr), a_arr) 
-- from cte_arr


-- ------------------------------------------------------------------------------------------------------
-- drop function if exists sm_sc.fv_d_redistr_zscore(float[], float[], int[]);
create or replace function sm_sc.fv_d_redistr_zscore
(
  i_depdt_var                  float[]                  ,                     -- zscore 输出
  i_indepdt_var                float[]                  ,                     -- zscore 算子的输入，来自上一层算子的输出
  i_cnt_per_grp                int[]
)
returns float[]
as
$$
declare 
  v_cur_y  int               ;
  v_cur_x  int               ;
  v_cur_x3 int               ;
  v_cur_x4 int               ;
  
begin
  if current_setting('pg4ml._v_is_debug_check', true) = '1'
  then
    if array_ndims(i_cnt_per_grp) > 1 
    then 
      raise exception 'unsupport ndims of i_cnt_per_grp > 1.';
    elsif array_ndims(i_indepdt_var) <> array_length(i_cnt_per_grp, 1)
    then 
      raise exception 'unmatch between ndims of i_indepdt_var and length of i_cnt_per_grp.';
    elsif array_ndims(i_depdt_var) > 4 or array_ndims(i_indepdt_var) > 4
      or array_dims(i_depdt_var) <> array_dims(i_indepdt_var) 
      -- -- or array_length(i_depdt_var, 1) <> array_length(i_indepdt_var, 1)
      -- -- or array_length(i_depdt_var, 2) <> array_length(i_indepdt_var, 2)
      -- -- or array_length(i_depdt_var, 3) <> array_length(i_indepdt_var, 3)
      -- -- or array_length(i_depdt_var, 4) <> array_length(i_indepdt_var, 4)
    then 
      raise exception 'unmatch ndims or length between i_depdt_var and i_indepdt_var.';
    elsif 
      0 <> any 
      (
        (
          select 
            array_agg(array_length(i_indepdt_var, a_cur_dim) order by a_cur_dim) 
          from generate_series(1, array_ndims(i_indepdt_var)) tb_a_cur_dim(a_cur_dim)
        )
        %` i_cnt_per_grp
      )
    then 
      raise exception 'unperfect i_indepdt_var''s length for i_cnt_per_grp at some dims';
    end if;
  end if;
  
  if i_indepdt_var is null
  then 
    return null;
  
  elsif i_cnt_per_grp is null 
  then 
    return sm_sc.fv_d_redistr_zscore(i_depdt_var, i_indepdt_var);
    
  elsif array_length(i_cnt_per_grp, 1) = 1
  then 
    for v_cur_y in 1 .. array_length(i_indepdt_var, 1) / i_cnt_per_grp[1]
    loop 
      i_depdt_var
        [(v_cur_y - 1) * i_cnt_per_grp[1] + 1 : v_cur_y * i_cnt_per_grp[1]] 
      := 
        sm_sc.fv_d_redistr_zscore
        (
          i_depdt_var
            [(v_cur_y - 1) * i_cnt_per_grp[1] + 1 : v_cur_y * i_cnt_per_grp[1]]
        , i_indepdt_var
            [(v_cur_y - 1) * i_cnt_per_grp[1] + 1 : v_cur_y * i_cnt_per_grp[1]]
        );
    end loop;
    return i_depdt_var;
    
  elsif array_length(i_cnt_per_grp, 1) = 2
  then    
    for v_cur_y in 1 .. array_length(i_indepdt_var, 1) / i_cnt_per_grp[1]
    loop 
      for v_cur_x in 1 .. array_length(i_indepdt_var, 2) / i_cnt_per_grp[2]
      loop 
        i_depdt_var
          [(v_cur_y - 1) * i_cnt_per_grp[1] + 1 : v_cur_y * i_cnt_per_grp[1]]
          [(v_cur_x - 1) * i_cnt_per_grp[2] + 1 : v_cur_x * i_cnt_per_grp[2]] 
        := 
          sm_sc.fv_d_redistr_zscore
          (
            i_depdt_var
              [(v_cur_y - 1) * i_cnt_per_grp[1] + 1 : v_cur_y * i_cnt_per_grp[1]]
              [(v_cur_x - 1) * i_cnt_per_grp[2] + 1 : v_cur_x * i_cnt_per_grp[2]]
          , i_indepdt_var
              [(v_cur_y - 1) * i_cnt_per_grp[1] + 1 : v_cur_y * i_cnt_per_grp[1]]
              [(v_cur_x - 1) * i_cnt_per_grp[2] + 1 : v_cur_x * i_cnt_per_grp[2]]
          );
      end loop;
    end loop;
    return i_depdt_var;
    
  elsif array_length(i_cnt_per_grp, 1) = 3
  then
    for v_cur_y in 1 .. array_length(i_indepdt_var, 1) / i_cnt_per_grp[1]
    loop 
      for v_cur_x in 1 .. array_length(i_indepdt_var, 2) / i_cnt_per_grp[2]
      loop 
        for v_cur_x3 in 1 .. array_length(i_indepdt_var, 3) / i_cnt_per_grp[3]
        loop 
          i_depdt_var
            [(v_cur_y - 1) * i_cnt_per_grp[1] + 1 : v_cur_y * i_cnt_per_grp[1]]
            [(v_cur_x - 1) * i_cnt_per_grp[2] + 1 : v_cur_x * i_cnt_per_grp[2]]
            [(v_cur_x3 - 1) * i_cnt_per_grp[3] + 1 : v_cur_x3 * i_cnt_per_grp[3]] 
          := 
            sm_sc.fv_d_redistr_zscore
            (
              i_depdt_var
                [(v_cur_y - 1) * i_cnt_per_grp[1] + 1 : v_cur_y * i_cnt_per_grp[1]]
                [(v_cur_x - 1) * i_cnt_per_grp[2] + 1 : v_cur_x * i_cnt_per_grp[2]]
                [(v_cur_x3 - 1) * i_cnt_per_grp[3] + 1 : v_cur_x3 * i_cnt_per_grp[3]]
            , i_indepdt_var
                [(v_cur_y - 1) * i_cnt_per_grp[1] + 1 : v_cur_y * i_cnt_per_grp[1]]
                [(v_cur_x - 1) * i_cnt_per_grp[2] + 1 : v_cur_x * i_cnt_per_grp[2]]
                [(v_cur_x3 - 1) * i_cnt_per_grp[3] + 1 : v_cur_x3 * i_cnt_per_grp[3]]
            );
        end loop;
      end loop;
    end loop;
    return i_depdt_var;
    
  elsif array_length(i_cnt_per_grp, 1) = 4
  then
    for v_cur_y in 1 .. array_length(i_indepdt_var, 1) / i_cnt_per_grp[1]
    loop 
      for v_cur_x in 1 .. array_length(i_indepdt_var, 2) / i_cnt_per_grp[2]
      loop 
        for v_cur_x3 in 1 .. array_length(i_indepdt_var, 3) / i_cnt_per_grp[3]
        loop 
          for v_cur_x4 in 1 .. array_length(i_indepdt_var, 4) / i_cnt_per_grp[4]
          loop 
            i_depdt_var
              [(v_cur_y - 1) * i_cnt_per_grp[1] + 1 : v_cur_y * i_cnt_per_grp[1]]
              [(v_cur_x - 1) * i_cnt_per_grp[2] + 1 : v_cur_x * i_cnt_per_grp[2]]
              [(v_cur_x3 - 1) * i_cnt_per_grp[3] + 1 : v_cur_x3 * i_cnt_per_grp[3]]
              [(v_cur_x4 - 1) * i_cnt_per_grp[4] + 1 : v_cur_x4 * i_cnt_per_grp[4]] 
            := 
              sm_sc.fv_d_redistr_zscore
              (
                i_depdt_var
                  [(v_cur_y - 1) * i_cnt_per_grp[1] + 1 : v_cur_y * i_cnt_per_grp[1]]
                  [(v_cur_x - 1) * i_cnt_per_grp[2] + 1 : v_cur_x * i_cnt_per_grp[2]]
                  [(v_cur_x3 - 1) * i_cnt_per_grp[3] + 1 : v_cur_x3 * i_cnt_per_grp[3]]
                  [(v_cur_x4 - 1) * i_cnt_per_grp[4] + 1 : v_cur_x4 * i_cnt_per_grp[4]]
              , i_indepdt_var
                  [(v_cur_y - 1) * i_cnt_per_grp[1] + 1 : v_cur_y * i_cnt_per_grp[1]]
                  [(v_cur_x - 1) * i_cnt_per_grp[2] + 1 : v_cur_x * i_cnt_per_grp[2]]
                  [(v_cur_x3 - 1) * i_cnt_per_grp[3] + 1 : v_cur_x3 * i_cnt_per_grp[3]]
                  [(v_cur_x4 - 1) * i_cnt_per_grp[4] + 1 : v_cur_x4 * i_cnt_per_grp[4]]
              );
          end loop;
        end loop;
      end loop;
    end loop;
    return i_depdt_var;
  end if;
end
$$
language plpgsql stable
parallel safe
cost 100;

-- with 
-- cte_arr as 
-- (
--   select 
--     sm_sc.fv_new_rand(array[4, 6, 9]) as a_arr
-- )
-- select 
--   sm_sc.fv_d_redistr_zscore
--   (
--     sm_sc.fv_redistr_zscore(a_arr)
--   , a_arr
--   , array[2, 3, 3]
--   ) 
-- from cte_arr